import numpy as np
import matplotlib.pyplot as plt
from pykalman import KalmanFilter

# 设置随机数种子，确保结果可复现
np.random.seed(42)

# 生成观测数据 - 带有噪声的正弦波
n_timesteps = 100
time = np.linspace(0, 3 * np.pi, n_timesteps)
true_signal = 20 * np.sin(time)
noise = 10 * np.random.randn(n_timesteps)  # 增加噪声幅度以便更明显地观察滤波效果
observations = true_signal + noise

# 配置卡尔曼滤波器
# 状态转移矩阵: [位置, 速度] -> [位置+速度, 速度]
transition_matrix = np.array([[1, 1], [0, 1]])

# 观测矩阵: 只能观测到位置
observation_matrix = np.array([[1, 0]])

# 转移协方差 - 控制状态转移的不确定性
transition_covariance = 0.01 * np.eye(2)

# 观测协方差 - 控制观测噪声的不确定性
observation_covariance = 10.0 * np.eye(1)  # 根据噪声幅度调整

# 初始状态和协方差
initial_state_mean = [0, 1]  # 初始位置和速度的估计
initial_state_covariance = np.eye(2)

# 创建卡尔曼滤波器实例
kf = KalmanFilter(
    transition_matrices=transition_matrix,
    observation_matrices=observation_matrix,
    transition_covariance=transition_covariance,
    observation_covariance=observation_covariance,
    initial_state_mean=initial_state_mean,
    initial_state_covariance=initial_state_covariance
)

# 使用期望最大化(EM)算法优化参数并进行平滑
kf = kf.em(observations, n_iter=5)
filtered_state_means, filtered_state_covariances = kf.filter(observations)
smoothed_state_means, smoothed_state_covariances = kf.smooth(observations)

# 预测未来5个时间步
n_future_steps = 5
future_times = np.linspace(time[-1], time[-1] + n_future_steps, n_future_steps)
future_states = np.zeros((n_future_steps, 2))
current_state = smoothed_state_means[-1]
current_covariance = smoothed_state_covariances[-1]

for i in range(n_future_steps):
    current_state = np.dot(transition_matrix, current_state)
    current_covariance = (np.dot(transition_matrix, np.dot(current_covariance, transition_matrix.T)) +
                         transition_covariance)
    future_states[i] = current_state

# 计算均方误差
filtered_mse = np.mean((filtered_state_means[:, 0] - true_signal) ** 2)
smoothed_mse = np.mean((smoothed_state_means[:, 0] - true_signal) ** 2)

# 可视化结果
plt.figure(figsize=(14, 10))

# 位置估计图
plt.subplot(2, 1, 1)
plt.scatter(time, observations, marker='x', color='blue', alpha=0.5, label='show')
plt.plot(time, true_signal, 'k-', label='ture')
plt.plot(time, filtered_state_means[:, 0], 'r--', label='after (MSE: {:.2f})'.format(filtered_mse))
plt.plot(time, smoothed_state_means[:, 0], 'g-', label='hua (MSE: {:.2f})'.format(smoothed_mse))
plt.plot(future_times, future_states[:, 0], 'm--', label='cai')
plt.fill_between(
    time,
    smoothed_state_means[:, 0] - 2 * np.sqrt(smoothed_state_covariances[:, 0, 0]),
    smoothed_state_means[:, 0] + 2 * np.sqrt(smoothed_state_covariances[:, 0, 0]),
    color='g', alpha=0.2, label='95%置信区间'
)
plt.title('kalman')
plt.xlabel('time')
plt.ylabel('ooo')
plt.legend()
plt.grid(True)

# 速度估计图
plt.subplot(2, 1, 2)
plt.plot(time, filtered_state_means[:, 1], 'r--', label='滤波后速度')
plt.plot(time, smoothed_state_means[:, 1], 'g-', label='平滑后速度')
plt.plot(future_times, future_states[:, 1], 'm--', label='预测速度')
plt.fill_between(
    time,
    smoothed_state_means[:, 1] - 2 * np.sqrt(smoothed_state_covariances[:, 1, 1]),
    smoothed_state_means[:, 1] + 2 * np.sqrt(smoothed_state_covariances[:, 1, 1]),
    color='g', alpha=0.2, label='95%置信区间'
)
plt.title('卡尔曼滤波器: 速度估计')
plt.xlabel('时间')
plt.ylabel('速度')
plt.legend()
plt.grid(True)

plt.tight_layout()
plt.show()

# 打印优化后的模型参数
print("优化后的转移矩阵:\n", kf.transition_matrices)
print("\n优化后的观测矩阵:\n", kf.observation_matrices)
print("\n优化后的转移协方差:\n", kf.transition_covariance)
print("\n优化后的观测协方差:\n", kf.observation_covariance)
print("\n初始状态均值:\n", kf.initial_state_mean)
print("\n初始状态协方差:\n", kf.initial_state_covariance)
